from asyncore import write
from dis import disco
from hashlib import new
from operator import ne
import torch
import torch.nn as nn
import numpy as np
from dqn import DQN
from focal_loss_example import FocalLoss
from model import *
from prepare_data import *
from utils_rename import evaluate
import torch.optim as optim
import random

from CNN import CNN_model
from resnet import ResNet, BasicBlock
from alexnet import AlexNet
from densenet import DenseNet
from vgg import VGG, make_layers
from squeezenet import SqueezeNet
from mixmim import MixMIMTransformer
from van import VAN
from Actor_Critic import ActorCritic
from active_learning_baseline import BadgeSampling

class RL_model(nn.Module):
    def __init__(self, opt, embedding_dim):
        super(RL_model, self).__init__()
        self.opt = opt
        self.visit_count = 0
        self.singleCapacity = opt.singleCapacity
        self.epochs = opt.epochs
        _, self.dimension_adj_martix = get_all_reflect_relation(opt.data_path)
        self.aes = AesModel(embedding_dim, opt).cuda()
        if opt.model_name == "CNN":
            self.aes = CNN_model(opt.Kernel_size1, opt.Stride1, opt.Kernel_size2, opt.Stride2).cuda()
        if opt.model_name == "ResNet":
            self.aes = ResNet(BasicBlock, [2, 2, 2, 2], num_classes=77).cuda()
        if opt.model_name == "AlexNet":
            self.aes = AlexNet(num_classes=77).cuda()
        if opt.model_name == "DenseNet":
            self.aes = DenseNet(num_init_features=64, growth_rate=32, block_config=(6, 12, 24, 16), num_classes=77).cuda()
        if opt.model_name == "VGG":
            self.aes = VGG(make_layers([64, 'M', 128, 'M', 256, 256, 'M', 512, 512, 'M', 512, 512, 'M']),
                        num_classes=77).cuda()
        if opt.model_name == "SqueezeNet":
            self.aesl = SqueezeNet(version=1, num_classes=77).cuda()
            # model = FusedGoogLeNet(flayer_cprate=[0.1]*64).cuda()
            # model = VisionTransformer(image_size=512, patch_size=opt.batchSize, embedding_dim=3 * 512 * 512, num_classes=77).cuda()
        if opt.model_name == "MixMIMTransformer":
            self.aes = MixMIMTransformer(img_size=opt.photo_size, patch_size=opt.batchSize, window_size=[16, 16, 16, 8]).cuda()
        if opt.model_name == "VAN":
            self.aes = VAN(arch='tiny', patch_sizes=[16, 8, 8, 8]).cuda() 
        if opt.model_name == "Badge":
            self.al_models = [BadgeSampling(opt) for _ in range(opt.num_classes)]
        
        self.aes_loss_func1 = FocalLoss(opt.FocalLoss_alpha, opt.FocalLoss_gamma)
        self.aes_loss_func2 = nn.BCEWithLogitsLoss()
        self.aes_loss_func3 = nn.L1Loss()
        # self.aes_optimizer = [optim.Adam(model.parameters(), lr=opt.lr) for model in self.aes.submodels]
        if opt.model_name == str('Aes') or opt.model_name == str('Actor_Critic') or opt.model_name == str('Random') or opt.model_name == str('Badge'):
            self.aes_cnn_optimizer = optim.Adam(self.aes.CNN.parameters(), lr=opt.lr)
            self.aes_submodels_optimizer = optim.Adam(self.aes.submodels.parameters(), lr=opt.lr)
        else:
            self.baseline_model_optimizer = optim.Adam(self.aes.parameters(), lr=opt.lr)

        self.dqn_state_nums = 7 + opt.embedding_dim
        if opt.state_type == str('1'):
            self.dqn_state_nums = 7
        elif opt.state_type == str('2'):
            self.dqn_state_nums = opt.embedding_dim
        elif opt.state_type ==  str('3'):
            self.dqn_state_nums = 77
        if self.opt.model_name == str('Actor_Critic'):
            self.dqnlist = [ActorCritic(self.dqn_state_nums, opt.hidden_dim, 2, opt.lr, opt.gamma, opt.epsilon,
                           opt.sample_size).cuda() for _ in range(opt.agent_nums)]
        elif self.opt.model_name == str('Aes'):
            self.dqnlist = [DQN(self.dqn_state_nums, opt.hidden_dim, 2, opt.lr, opt.gamma, opt.epsilon, opt.target_update_nums,
                         opt.ReplayBuffer_capacity, opt.sample_size).cuda() for _ in range(opt.agent_nums)]
        self.train_agents, self.test_agents = {}, {}
        if opt.Save_dimension_chosen_nums == str('true'):
            self.dimension_chosen_nums = [0 for _ in range(opt.num_classes)]
            self.dimension_abs_dif_sum = [0 for _ in range(opt.num_classes)]
            print('Save_dimension_chosen_nums true')
        self.ReplayBuffer_capacity = opt.ReplayBuffer_capacity

        self.state_embedding = nn.Embedding(opt.agent_nums, opt.embedding_dim).cuda()
        self.state_embed = self.state_embedding(torch.LongTensor([_ for _ in range(opt.agent_nums)]).cuda()).cuda()
        self.GAT = GAT(self.opt.embedding_dim, self.opt.embedding_dim, self.opt.embedding_dim, 0.7, 0.2, self.opt.num_heads).cuda()


    def cal_s1(self, chosen_dimensions_val):
        std_val, avg_val, min_val, max_val = torch.std(chosen_dimensions_val), torch.mean(chosen_dimensions_val), torch.min(chosen_dimensions_val), torch.max(chosen_dimensions_val)
        m1, m2, m3 = torch.quantile(chosen_dimensions_val, 0.25), torch.quantile(chosen_dimensions_val, 0.50), torch.quantile(chosen_dimensions_val, 0.75)
        return torch.tensor([std_val, avg_val, min_val, max_val, m1, m2, m3], dtype=torch.float).unsqueeze(0).cuda()

    def cal_s2(self, mark_id_list, state_embed):
        s2 = self.state_embed[mark_id_list]
        s2 = torch.mean(s2, dim=0).unsqueeze(0).cuda()
        return s2

    def cal_new_state(self, action_list, aes_ability):
        mark_id_list = []
        for _ in range(len(action_list)):
            if action_list[_] == 1:
                mark_id_list.append(_)
        mark_id_list = torch.LongTensor(mark_id_list)
        s1 = []
        if len(mark_id_list) == 0:
            s1 = torch.tensor([0 for _ in range(7)]).unsqueeze(0).cuda()
        else:
            chosen_dimensions_val = aes_ability[:, mark_id_list]
            s1 = self.cal_s1(chosen_dimensions_val)
        s2 = self.cal_s2(mark_id_list, self.state_embed)
        new_state = torch.cat((s1, s2), dim=1).cuda()
        if self.opt.state_type == str('1'):
            new_state = s1.cuda().float()
        elif self.opt.state_type == str('2'):
            new_state = s2.cuda().float()
        elif self.opt.state_type == str('3'):
            new_state = aes_ability
        return new_state

    def init_state(self, aes_ability):
        action_list = [random.randint(0, 1) for _ in range(self.opt.agent_nums)]
        ini_state = self.cal_new_state(action_list, aes_ability)
        return ini_state, action_list

    def cal_new_ability(self, action_list, aes_ability, mark_ability):
        aes_ability, mark_ability, action_list = aes_ability.tolist()[0], mark_ability.tolist()[0], torch.tensor(action_list)
        cur_ability = torch.where(action_list==1, torch.tensor(mark_ability), torch.tensor(aes_ability))

        return torch.tensor(cur_ability).unsqueeze(0).cuda()

    def pearson_correlation(self, martix):
        martix = martix.cpu().detach().numpy()
        correlation_matrix = np.corrcoef(martix)

        # 提取上三角部分的相关性（不包括对角线）
        upper_triangle_indices = np.triu_indices_from(correlation_matrix, k=1)
        upper_triangle_values = correlation_matrix[upper_triangle_indices]

        # 计算平均相关性值
        average_correlation = np.mean(upper_triangle_values)
        return average_correlation

    def cal_reward(self, last_mse, last_mae, action_list, mark_ability, new_ability, discount):
        mae, mse, rmse, r2, pea, mape, mpe = evaluate(new_ability.tolist(), mark_ability.tolist())
        mse_improve = (last_mse - mse) * self.opt.msepara
        mae_improve = (last_mae - mae) * self.opt.maepara
        mark_id_list = []
        for _ in range(len(action_list)):
            if action_list[_] == 1:
                mark_id_list.append(_)
        sim = 0
        if len(mark_id_list) < 2:
            sim = 10
        else:
            sim_martix = self.state_embed[torch.tensor(mark_id_list)].cuda()
            sim = self.pearson_correlation(sim_martix)
        punish_val = sum(action_list) * self.opt.mu

        reward = (mse_improve + mae_improve + sim - punish_val) * discount
        if self.opt.reward_type == str('1'):
            reward = (mse_improve - punish_val) * discount
        elif self.opt.reward_type == str('2'):
            reward = (mae_improve - punish_val) * discount
        elif self.opt.reward_type == str('3'):
            reward = (sim - punish_val) * discount
        elif self.opt.reward_type == str('4'):
            reward = (mse_improve + mae_improve - punish_val) * discount
        record_list = [mse_improve, sim, punish_val]
        return reward, record_list

    def RL_train_Aes_Submodels(self, img, labels):
        scores, loss_vector = torch.zeros(1,self.opt.num_classes).cuda(), torch.zeros(1, self.opt.num_classes).cuda()
        score = self.aes(img).cuda()
        min_val = score.min()
        max_val = score.max()

        # min-max 归一化到 [0.2, 0.8]
        # score = (score - min_val) * (0.8 - 0.2) / (max_val - min_val) + 0.2
        score = (score - min_val) / (max_val - min_val)


        loss1 = [self.aes_loss_func1(score[0][_].cuda(), labels[_].cuda()) for _ in range(self.opt.num_classes)]
        loss2 = [self.aes_loss_func2(score[0][_].cuda(), labels[_].cuda()) for _ in range(self.opt.num_classes)]
        loss3 = [self.aes_loss_func3(score[0][_].cuda(), labels[_].cuda()) for _ in range(self.opt.num_classes)]
        loss_vector = [loss1[_] + loss2[_] + loss3[_] for _ in range(len(loss1))]
        scores[0] = score

        return scores, loss_vector

    def train_Aes_CNN(self, train_loader):
        self.aes.train()
        for img, labels in train_loader:
            labels = labels.unsqueeze(0).cuda()
            img = img.unsqueeze(0).cuda()
            score = self.aes(img).cuda()
            min_val = score.min()
            max_val = score.max()

            # min-max 归一化到 [0.2, 0.8]
            # score = (score - min_val) * (0.8 - 0.2) / (max_val - min_val) + 0.2
            score = (score - min_val) / (max_val - min_val)

            loss1 = self.aes_loss_func1(score, torch.tensor(labels).cuda())
            loss2 = self.aes_loss_func2(score, torch.tensor(labels).cuda())
            loss3 = self.aes_loss_func3(score, torch.tensor(labels).cuda())
            loss = loss1 + loss2 + loss3
            if self.opt.model_name == str('Aes') or self.opt.model_name == str('Actor_Critic') or self.opt.model_name == str('Random') or self.opt.model_name == str('Badge'):
                self.aes_cnn_optimizer.zero_grad()
            else:
                self.baseline_model_optimizer.zero_grad()
            loss.backward()
            if self.opt.model_name == str('Aes') or self.opt.model_name == str('Actor_Critic') or self.opt.model_name == str('Random') or self.opt.model_name == str('Badge'):
                self.aes_cnn_optimizer.step()
            else:
                self.baseline_model_optimizer.step()

    def test(self, test_loader):
        import time
        st = time.time()
        true_list, rec_list = [], []
        # with torch.inference_mode():
        for img, name, label in test_loader:
            if len(img) < self.opt.batchSize:
                continue
            img = img.cuda()
            score = self.aes(img).cuda()

            min_val = score.min()
            max_val = score.max()
            # min-max 归一化到 [0.2, 0.8]
            # score = (score - min_val) * (0.8 - 0.2) / (max_val - min_val) + 0.2
            score = (score - min_val) / (max_val - min_val)

            rec_list.append(score.tolist())
            true_list.append(label.tolist())
        en = time.time()
        print(f'test cost time {en - st}')
        return evaluate(rec_list, true_list)

    def train_RL(self, train_loader, test_loader):
        num_visits, tl_visits = 0, self.singleCapacity * self.opt.path_len
        self.aes.train()
        self.train_agents = {}
        print('tl_visits are', tl_visits)
        while num_visits < tl_visits:
            print('num_visits are', num_visits)
            action_martix, loss_matrix, rewards, CNN_Loss, submodels_loss = [], [], [], [], []
            for img, name, labels in train_loader:
                if name not in self.train_agents:
                    self.train_agents[name] = {}
                    one_picture = img.unsqueeze(0).cuda()

                    # aes_ability = torch.tensor([model(one_picture) for model in self.aes]).unsqueeze(0).cuda()
                    aes_ability = self.aes(one_picture).cuda()

                    min_val = aes_ability.min()
                    max_val = aes_ability.max()

                    # min-max 归一化到 [0.2, 0.8]
                    # aes_ability = (aes_ability - min_val) * (0.8 - 0.2) / (max_val - min_val) + 0.2
                    aes_ability = (aes_ability - min_val) / (max_val - min_val)

                    self.train_agents[name]['cur_state'], self.train_agents[name]['last_actions'] = self.init_state(aes_ability)
                    self.train_agents[name]['mark_ability'] = torch.tensor(labels).unsqueeze(0).cuda()
                    self.train_agents[name]['new_ability'] = labels

                    self.train_agents[name]['aes_ability'] = aes_ability
                    score, label = aes_ability.tolist(), labels.unsqueeze(0).tolist()
                    mae, mse, rmse, r2, pea, mape, mpe = evaluate(score, label)

                    self.train_agents[name]['last_mse'] = mse
                    self.train_agents[name]['last_mae'] = mae
                    self.train_agents[name]['discount'] = self.opt.gamma
                    self.train_agents[name]['iter_count'] = 0
                    self.train_agents[name]['rewards'] = []

                if self.train_agents[name]['iter_count'] >= self.opt.path_len:
                    continue
                if self.opt.model_name == str('Random') or self.opt.model_name == str('Badge'):
                    action_list = [1] * 8 + [0] * (self.opt.num_classes - 8)
                    random.shuffle(action_list)
                else:
                    action_list = [dqn.choose_action(self.train_agents[name]['cur_state']) for dqn in self.dqnlist]
                if self.opt.Save_dimension_chosen_nums == str('true'):
                    predict_score, _ = self.RL_train_Aes_Submodels(img.unsqueeze(0).cuda(), labels)
                    for idx in range(len(action_list)):
                        if action_list[idx] == 1:
                            self.dimension_chosen_nums[idx] += 1
                            self.dimension_abs_dif_sum[idx] += abs(predict_score[0][idx] - labels[idx])
                # action_martix.append(torch.tensor(action_list).unsqueeze(0).cuda())
                action_martix.append(torch.tensor(action_list).unsqueeze(0).cuda())

                if self.opt.model_name != str('Random') and self.opt.model_name != str('Badge'):
                    new_state = self.cal_new_state(action_martix[-1].squeeze(0).tolist(), self.train_agents[name]['aes_ability'])
                    self.train_agents[name]['new_state'] = new_state
                    # self.train_agents[name]['last_actions'] = action_list
                self.train_agents[name]['iter_count'] += 1
                num_visits += 1
            
            print("time end")
            action_martix = torch.cat(action_martix, dim=0).cuda()

            if self.opt.model_name == str('Badge'):
                # print(f'action_martix is {action_martix}')
                action_martix = torch.zeros_like(action_martix)
                X, Y = [_[0].unsqueeze(0) for _ in train_loader], [_[-1].unsqueeze(0) for _ in train_loader]
                X, Y = torch.cat(X, dim=0).cuda(), torch.cat(Y, dim=0).cuda()
                for idx in range(self.opt.num_classes):
                    output = self.al_models[idx].query(self.opt.query_num, X, Y)
                    action_martix[output, idx] = 1
                    # print(f'idx is {idx}, output is {output}')
                    self.al_models[idx].train(X, Y)
                # print(f'action_martix2 is {action_martix}')


            for epoch in range(self.opt.model_epochs):
                new_ability_martix, loss_matrix = [], []
                for img, name, labels in train_loader:
                    new_ability_vector, loss_vector = self.RL_train_Aes_Submodels(img.unsqueeze(0).cuda(), labels)
                    new_ability_martix.append(new_ability_vector)
                    loss_matrix.append(loss_vector)

                loss_matrix_tensor = torch.stack([torch.stack(row) for row in loss_matrix]).cuda()
                result_matrix = torch.mm(loss_matrix_tensor.T.cuda(), torch.tensor(action_martix, dtype=torch.float32, requires_grad=True).cuda())
                total_loss = result_matrix.sum()
                if self.opt.model_name == str('Aes') or self.opt.model_name == str('Actor_Critic') or self.opt.model_name == str('Random') or self.opt.model_name == str('Badge'):
                    self.aes_cnn_optimizer.zero_grad()
                else:
                    self.baseline_model_optimizer.zero_grad()
                total_loss.backward()
                CNN_Loss.append(total_loss.item())
                if self.opt.model_name == str('Aes') or self.opt.model_name == str('Actor_Critic') or self.opt.model_name == str('Random') or self.opt.model_name == str('Badge'):
                    self.aes_cnn_optimizer.step()
                else:
                    self.baseline_model_optimizer.step()

                if self.opt.model_name == str('Aes') or self.opt.model_name == str('Actor_Critic') or self.opt.model_name == str('Random') or self.opt.model_name == str('Badge'):
                    new_ability_martix, loss_matrix = [], []
                    for img, name, labels in train_loader:
                        new_ability_vector, loss_vector = self.RL_train_Aes_Submodels(img.unsqueeze(0).cuda(), labels)
                        new_ability_martix.append(new_ability_vector)
                        loss_matrix.append(loss_vector)
                    
                    loss_matrix_tensor = torch.stack([torch.stack(row) for row in loss_matrix]).cuda()
                    result_matrix = torch.mm(loss_matrix_tensor.T.cuda(), torch.tensor(action_martix, dtype=torch.float32, requires_grad=True).cuda())
                    total_loss = result_matrix.sum()
                    col_loss = torch.sum(result_matrix, dim=0)
                    self.aes_submodels_optimizer.zero_grad()
                    total_loss.backward(retain_graph=True)
                    one_sub_loss = col_loss.tolist()
                    submodels_loss.append(one_sub_loss)
                    self.aes_submodels_optimizer.step()

                if epoch == self.opt.model_epochs - 1:
                    if self.opt.model_name != str('Random') and self.opt.model_name != str('Badge'):
                        new_ability_martix = torch.cat(new_ability_martix, dim=0).cuda()
                        for (img, name, labels), action_list, new_ability in zip(train_loader, action_martix, new_ability_martix):
                            reward, record_list = self.cal_reward(self.train_agents[name]['last_mse'], self.train_agents[name]['last_mae'], action_list.tolist(), self.train_agents[name]['mark_ability'], new_ability.unsqueeze(0), self.train_agents[name]['discount'])
                            for idx, dqn in enumerate(self.dqnlist):
                                if action_list[idx].item() == 1 or self.train_agents[name]['last_actions'][idx] == 1:
                                    if self.opt.model_name == str('Actor_Critic'):
                                        dqn.update(self.train_agents[name]['cur_state'], reward.item(), self.train_agents[name]['new_state'])
                                    else:
                                        dqn.store_transition(torch.squeeze(self.train_agents[name]['cur_state'], dim=0).tolist(), action_list[idx].item(), reward.item(), torch.squeeze(self.train_agents[name]['new_state'], dim=0).tolist())
                            self.train_agents[name]['last_actions'] = action_list.tolist()
                            if self.opt.model_name != str('Actor_Critic'):
                                for dqn in self.dqnlist:
                                    if dqn.memery.size() >= self.opt.min_size:
                                        dqn.update()

                        self.train_agents[name]['cur_state'], self.train_agents[name]['discount'] = self.train_agents[name]['new_state'], self.train_agents[name]['discount'] * self.opt.gamma
            mae, mse, rmse, r2, pea, mape, mpe = self.test(test_loader)
            avg_num = torch.sum(action_martix) / self.singleCapacity
            self.train_agents[name]['last_mse'], self.train_agents[name]['last_mae'] = mse, mae
            mae, mse, rmse, r2, pea, mape, mpe, avg_num = [metric.item() for metric in (mae, mse, rmse, r2, pea, mape, mpe, avg_num)]

            print(f'{mae} {mse} {rmse} {r2} {pea} {mape} {mpe}')
            with open(f"/home/n40/code/aesthetic_path/AL_MRL/result/{self.opt.expand_name}.txt", mode='a', encoding='utf-8') as file:
                ans = str(mae) + str(' ') + str(mse) + str(' ') + str(rmse) + str(' ') + str(r2) + str(' ') + str(pea) + str(' ') + str(mape) + str(' ') + str(mpe) + str(' ') + str(avg_num) + str('\n')
                file.write(ans)
            with open(f"/home/n40/code/aesthetic_path/AL_MRL/loss/{self.opt.expand_name}.txt", mode='a', encoding='utf-8') as file:
                for idx in range(len(CNN_Loss)):
                    ans = str(CNN_Loss[idx])
                    if self.opt.model_name == str('Aes') or self.opt.model_name == str('Actor_Critic'):
                        ans = ans + str(' ') + str(submodels_loss[idx]) + str('\n')
                    ans = ans + str('\n')
                    file.write(ans)
                file.write('\n')

            self.state_embed = self.GAT(self.state_embed, self.dimension_adj_martix).cuda()

    def compare_evaluate_result(self, left, right):
        lmae, lmse, lrmse, rmae, rmse, rrmse = left[0], left[1], left[2], right[0], right[1], right[2]
        if lmse < rmse:
            if lmae < rmae:
                return -1
            else:
                return -1 if (rmse - lmse >= lmae - rmae) else 1
        else:
            if lmae > rmae:
                return 1
            else:
                return 1 if (lmse - rmse >= rmae - lmae) else -1


    def fit(self, train_loader, test_loader):
        self.train_agents, trainingPictureData, trainedRlData, trainingRlData, data_set, numberOfRlImages, tl_random_data = {}, [], [], [], set(), 1000000, []
        name_to_data_dict, cur_data_idx, sorted_indices = {}, 0, []

        if self.opt.Set_data == str('true'):
            for img, name, labels in train_loader:
                for i in range(len(img)):
                    name_to_data_dict[name[i]] = [img[i], name[i], labels[i]]
            for img, name, labels in test_loader:
                for i in range(len(img)):
                    name_to_data_dict[name[i]] = [img[i], name[i], labels[i]]

            names_list = []
            with open(f'random_data.txt', mode='r', encoding='utf-8') as file:
                lines = file.readlines()
                one_name_list = []
                for line in lines:
                    row = line.strip()
                    if not row:
                        names_list.append(one_name_list)
                        one_name_list = []
                    else:
                        one_name_list.append(row)
            ans_list = []
            with open(f'/home/n40/code/aesthetic_path/AL_MRL/result/CNN_USE_TRANSFORM_new.txt', mode='r', encoding='utf-8') as file:
                lines = file.readlines()
                result_list = [100 for _ in range(8)]
                for line in lines:
                    row = line.strip().split()
                    if len(row) == 8:
                        numbers = [round(float(num), 4) for num in row]
                        if self.compare_evaluate_result(numbers, result_list) < 0:
                            result_list = numbers
                    else:
                        ans_list.append(result_list)
                        result_list = [100 for _ in range(7)]


            from functools import cmp_to_key
            def compare_indexed_results(x, y):
                return self.compare_evaluate_result(x[1], y[1])
            indexed_ans_list = list(enumerate(ans_list))
            sorted_indexed_ans_list = sorted(indexed_ans_list, key=cmp_to_key(compare_indexed_results))

            sorted_indices = [item[0] for item in sorted_indexed_ans_list][::-1]

        if self.opt.Save_data == str('true'):
            with open(f"/home/n40/code/aesthetic_path/AL_MRL/random_data.txt", mode='a', encoding='utf-8') as file:
                for one_sampleData in tl_random_data:
                    for img, name, labels in one_sampleData:
                        file.write(str(name))
                        file.write('\n')
                    file.write('\n')

        for img, name, labels in train_loader:
            for i in range(len(img)):
                data_set.add(name[i])
                if len(data_set) < self.singleCapacity:
                    trainingPictureData.append([img[i], labels[i]])
                elif len(data_set) == self.singleCapacity:
                    for epoch in range(self.opt.model_epochs):
                        self.train_Aes_CNN(trainingPictureData)
                    numberOfRlImages = self.singleCapacity
                elif numberOfRlImages != 0:
                    trainingRlData.append([img[i], name[i], labels[i]])
                    numberOfRlImages -= 1
                else:
                    sampledData = random.sample(trainedRlData, self.singleCapacity - len(trainingRlData)) + trainingRlData
                    if self.opt.Set_data == str('true'):
                        sampledData = []        
                        for name1 in names_list[sorted_indices[cur_data_idx]]:
                            sampledData.append(name_to_data_dict[name1])
                        cur_data_idx += 1


                    tl_random_data.append(sampledData)
                    self.train_RL(sampledData, test_loader)
                    # with open(f"result/{self.opt.expand_name}.txt", mode='a', encoding='utf-8') as file:
                    #     file.write('one batch end\n')
                    with open(f"/home/n40/code/aesthetic_path/AL_MRL/loss/{self.opt.expand_name}.txt", mode='a', encoding='utf-8') as file:
                        file.write('one batch end\n')
                    trainedRlData =  trainedRlData + trainingRlData
                    numberOfRlImages, trainingRlData = int(self.singleCapacity * 0.8), []
        if self.opt.is_adjust_parameter == str('true'):
            if hasattr(self.opt, self.opt.expand_name):
                adjust_parameter_val = getattr(self.opt, self.opt.expand_name)
            else:
                adjust_parameter_val = 'none'
            result_list = [100 for _ in range(7)]
            with open(f"/home/n40/code/aesthetic_path/AL_MRL/result/{self.opt.expand_name}.txt", mode='r', encoding='utf-8') as file:
                lines = file.readlines()
                for line in lines:
                    numbers = [float(num) for num in line.split()]
                    if self.compare_evaluate_result(numbers, result_list) < 0:
                        result_list = numbers            
            with open(f"/home/n40/code/aesthetic_path/AL_MRL/parameter_{self.opt.expand_name}.txt", mode='a', encoding='utf-8') as file:
                file.write(str(f'{self.opt.expand_name} {adjust_parameter_val} {result_list}'))
                file.write('\n')

        if self.opt.Save_dimension_chosen_nums == str('true'):
            with open(f"/home/n40/code/aesthetic_path/AL_MRL/dimension_chosen_nums.txt", mode='a', encoding='utf-8') as file:
                for idx in range(self.opt.num_classes):
                    file.write(str(f'{idx} {self.dimension_chosen_nums[idx]} {self.dimension_abs_dif_sum[idx]}'))
                    file.write('\n')
